package com.hcc.repository.extension.interceptor.pagination.dialect;

import com.hcc.repository.core.utils.CollUtils;
import com.hcc.repository.extension.interceptor.pagination.PaginationContext;
import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Function;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.Distinct;
import net.sf.jsqlparser.statement.select.GroupByElement;
import net.sf.jsqlparser.statement.select.Join;
import net.sf.jsqlparser.statement.select.OrderByElement;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectExpressionItem;
import net.sf.jsqlparser.statement.select.SelectItem;

import java.util.Collections;
import java.util.List;
import java.util.Optional;

/**
 * 抽象分页处理器
 *
 * @author hushengjun
 * @date 2023/5/1
 */
public abstract class AbstractDialect implements IDialect {

    @Override
    public void handle(PaginationContext context) {
        this.handlePageSql(context);
        if (!context.isNeedCount()) {
            return;
        }
        this.handleCountSql(context);
    }

    /**
     * 分页语句
     * @param context
     */
    protected abstract void handlePageSql(PaginationContext context);

    /**
     * count语句，简单的子查询
     * @param originalSql
     */
    private String simpleCountSql(String originalSql) {
        return "SELECT COUNT(*) FROM ( " + originalSql + " ) TOTAL";
    }

    /**
     * 优化的count语句，参考mybatis-plus的实现
     * @param originalSql
     * @return
     */
    protected String optimizeCountSql(String originalSql) {
        try {
            Select select = (Select) CCJSqlParserUtil.parse(originalSql);
            PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
            Distinct distinct = plainSelect.getDistinct();
            GroupByElement groupBy = plainSelect.getGroupBy();
            List<OrderByElement> orderBys = plainSelect.getOrderByElements();

            // 去除ORDER BY
            if (CollUtils.isNotEmpty(orderBys)) {
                plainSelect.setOrderByElements(null);
            }

            // select中含有问号则无需优化
            for (SelectItem selectItem : plainSelect.getSelectItems()) {
                if (selectItem.toString().contains("?")) {
                    return this.simpleCountSql(select.toString());
                }
            }

            // 若sql中含有GROUP BY或DISTINCT，则无需优化
            if (distinct != null || groupBy != null) {
                return this.simpleCountSql(select.toString());
            }

            // 判断是否可以去除join
            List<Join> joins = plainSelect.getJoins();
            if (CollUtils.isNotEmpty(joins)) {
                boolean canRemoveJoin = true;
                String whereStr = Optional.ofNullable(plainSelect.getWhere())
                        .map(Expression::toString)
                        .orElse("");
                // 不区分大小写
                whereStr = whereStr.toLowerCase();
                for (Join join : joins) {
                    if (!join.isLeft()) {
                        canRemoveJoin = false;
                        break;
                    }
                    Table table = (Table) join.getRightItem();
                    String str = Optional.ofNullable(table.getAlias()).map(Alias::getName).orElse(table.getName()) + ".";
                    // 不区分大小写
                    str = str.toLowerCase();
                    String onExpressionS = join.getOnExpression().toString();
                    // 如果 join 里包含 ?(代表有入参) 或者 where 条件里包含使用 join 的表的字段作条件,就不移除 join
                    if (onExpressionS.contains("?") || whereStr.contains(str)) {
                        canRemoveJoin = false;
                        break;
                    }
                }
                if (canRemoveJoin) {
                    plainSelect.setJoins(null);
                }
            }
            plainSelect.setSelectItems(this.countSelectItem());

            return select.toString();
        } catch (Exception e) {
            return this.simpleCountSql(originalSql);
        }
    }

    /**
     * COUNT(*)
     *
     * @return
     */
    private List<SelectItem> countSelectItem() {
        Function function = new Function();
        function.setName("COUNT");
        function.setAllColumns(true);
        return Collections.singletonList(new SelectExpressionItem(function));
    }

    /**
     * count语句
     * @param context
     */
    protected void handleCountSql(PaginationContext context) {
        String originalSql = context.getOriginalSql();
        if (context.isOptimizeCountSql()) {
            context.setCountSql(this.optimizeCountSql(originalSql));
            return;
        }
        context.setCountSql(this.simpleCountSql(originalSql));
    }

}
